import geoip2.database
import pandas as pd

from base.constant import EmptyValue
from base.util.common_util import get_project_root
from base.util.file_util import get_config
from base.util.geo_util import get_geo_by_ip, is_stable_name, fill_geo_empty


def replace_fields(df1, df2, mapping_dict):
    df1_extended = df1.sample(n=len(df2), replace=True, random_state=42).reset_index(drop=True)
    for df2_col, df1_col in mapping_dict.items():
        if df1_col in df1.columns and df2_col in df2.columns:
            df2[df2_col] = df1_extended[df1_col]
    return df2


if __name__ == '__main__':
    config_yaml, config_file = get_config()
    extract_sample = config_yaml['extract_sample']
    destination_file = pd.read_csv(f"{get_project_root()}{extract_sample['destination']}")
    final_file = pd.DataFrame()
    for i, path in enumerate(extract_sample['source']):
        print(f"[替换五元组]正在读取文件{i + 1}/{len(extract_sample['source'])}|{path}")
        file = pd.read_csv(f"{get_project_root()}{path}")
        final_file = pd.concat([file, final_file], axis=0, ignore_index=True)
    print("[替换五元组]开始")
    df = replace_fields(final_file, destination_file, extract_sample['mapping'])
    reader = geoip2.database.Reader(f'{get_project_root()}base/util/GeoLite2-City.mmdb')
    df = df.apply(lambda x: get_geo_by_ip(reader, x), axis=1)
    df.to_csv(f"{get_project_root()}{extract_sample['save_to']}", index=False)
    print("[替换五元组]替换完毕")
    config_file.close()
